import random
import numpy as np
import torch
from torchvision import datasets
from torchvision.transforms import transforms
from cnn_mnist.my_model import LeNet5
from cnn_mnist.utils import tool, plt_util


def init():
    # 设置随机种子，保证结果可复现
    torch.manual_seed(42)
    np.random.seed(42)
    random.seed(42)

    batch_size = 64
    learning_rate = 0.001

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # transform = transforms.Compose(
    #     [
    #         transforms.Resize((32, 32)),
    #         transforms.ToTensor(),  # 归一化 像素 /255,  h w c --> c h w
    #         transforms.Normalize(0.1307, 0.3081)
    #      ]
    # )

    transform = transforms.Compose([
        transforms.Pad(padding=50, fill=0),
        transforms.RandomRotation(15),  # 随机旋转
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # 加载数据集
    # 训练集
    train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)

    # 测试集
    test_dataset = datasets.MNIST('data', train=False, download=True, transform=transform)

    # 数据加载器
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

    model = LeNet5().to(device)

    # 损失函数
    criterion = torch.nn.CrossEntropyLoss()

    # 优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    return device, model, criterion, optimizer, train_loader, test_loader


def main():
    device, model, criterion, optimizer, train_loader, test_loader = init()

    epochs = 10
    train_losses = []
    train_accuracies = []

    test_losses = []
    test_accuracies = []

    for epoch in range(1, epochs + 1):
        train_loss, train_acc = tool.train(model, device, train_loader, optimizer, epoch, criterion)
        test_loss, test_acc = tool.test(model, device, test_loader, criterion)

        train_losses.append(train_loss)
        train_accuracies.append(train_acc)

        test_losses.append(test_loss)
        test_accuracies.append(test_acc)

    torch.save(model.state_dict(), 'lenet5_mnist_fully_sequential.pth')

    plt_util.plt_show(epochs, train_losses, test_losses, train_accuracies, test_accuracies)

    plt_util.visualize_predictions(model, device, test_loader)


if __name__ == '__main__':
    main()

